"""
图像净化处理系统 - V4.0零信任架构
防御Polyglot文件攻击和恶意载荷注入

安全特性:
- 强制图像重新编码净化
- 剥离所有非像素数据
- 多层恶意载荷检测
- EXIF和元数据清理
- 文件格式验证
"""

import io
import structlog
from typing import Tuple, Dict, Any, Optional
from pathlib import Path
from PIL import Image, ImageFile, UnidentifiedImageError
from PIL.ExifTags import TAGS
import magic
import re

logger = structlog.get_logger()


# 允许加载可能截断的图像（防止某些攻击）
ImageFile.LOAD_TRUNCATED_IMAGES = True


class ImageSanitizationError(Exception):
    """图像净化相关异常"""
    pass


class ImageSanitizer:
    """图像净化器 - 零信任图像处理"""

    def __init__(self):
        # 允许的图像格式
        self.allowed_formats = {'JPEG', 'PNG', 'BMP', 'RGB', 'RGBA'}

        # 恶意模式检测
        self.malicious_patterns = [
            r'<script[^>]*>.*?</script>',  # JavaScript
            r'<iframe[^>]*>.*?</iframe>',  # iframe
            r'eval\s*\(',  # eval函数
            r'exec\s*\(',  # exec函数
            r'system\s*\(',  # system调用
            r'shell_exec\s*\(',  # shell_exec
            r'base64_decode',  # base64解码
            r'php://',  # PHP包装器
            r'data://',  # 数据URI
            r'file://',  # 文件URI
        ]

        # 危险文件特征
        self.dangerous_signatures = [
            b'MZ\x90\x00',  # PE可执行文件
            b'\x7fELF',  # ELF可执行文件
            b'\xca\xfe\xba\xbe',  # Java类文件
            b'\xfe\xed\xfa\xce',  # Mach-O可执行文件
            b'PK\x03\x04',  # ZIP文件
            b'PK\x05\x06',  # ZIP文件
            b'PK\x07\x08',  # ZIP文件
        ]

    def sanitize_image(self, image_bytes: bytes, max_size_mb: int = 50) -> Tuple[bytes, Dict[str, Any]]:
        """
        净化图像数据

        Args:
            image_bytes: 原始图像字节数据
            max_size_mb: 最大允许文件大小（MB）

        Returns:
            (净化后的图像字节数据, 处理报告)

        Raises:
            ImageSanitizationError: 净化失败时抛出
        """
        try:
            logger.info("开始图像净化处理", input_size=len(image_bytes))

            # 初始化处理报告
            report = {
                'original_size': len(image_bytes),
                'sanitized_size': 0,
                'format_detected': None,
                'security_issues': [],
                'cleaning_actions': [],
                'is_safe': True,
                'processing_time': 0,
                'metadata_removed': {}
            }

            import time
            start_time = time.time()

            # 第一阶段：基础安全检查
            self._basic_security_check(image_bytes, report)

            # 第二阶段：文件大小验证
            self._validate_file_size(image_bytes, max_size_mb, report)

            # 第三阶段：文件头验证
            self._validate_file_header(image_bytes, report)

            # 第四阶段：恶意载荷扫描
            self._scan_for_malicious_payloads(image_bytes, report)

            # 第五阶段：图像加载和验证
            sanitized_image = self._load_and_validate_image(image_bytes, report)

            # 第六阶段：强制重新编码净化
            sanitized_bytes = self._force_reencoding(sanitized_image, report)

            # 第七阶段：最终验证
            self._final_validation(sanitized_bytes, report)

            report['sanitized_size'] = len(sanitized_bytes)
            report['processing_time'] = time.time() - start_time

            logger.info("图像净化完成",
                       original_size=report['original_size'],
                       sanitized_size=report['sanitized_size'],
                       security_issues=len(report['security_issues']),
                       is_safe=report['is_safe'])

            return sanitized_bytes, report

        except Exception as e:
            logger.error("图像净化过程失败", error=str(e))
            raise ImageSanitizationError(f"图像净化失败: {e}")

    def _basic_security_check(self, image_bytes: bytes, report: Dict[str, Any]) -> None:
        """基础安全检查"""
        # 检查是否为空文件
        if len(image_bytes) == 0:
            report['security_issues'].append("空文件")
            report['is_safe'] = False
            raise ImageSanitizationError("检测到空文件")

        # 检查文件大小是否异常小
        if len(image_bytes) < 100:
            report['security_issues'].append("文件过小，可能不是有效图像")
            report['is_safe'] = False
            raise ImageSanitizationError("文件过小，可能不是有效图像")

    def _validate_file_size(self, image_bytes: bytes, max_size_mb: int, report: Dict[str, Any]) -> None:
        """验证文件大小"""
        max_size_bytes = max_size_mb * 1024 * 1024
        if len(image_bytes) > max_size_bytes:
            report['security_issues'].append(f"文件过大: {len(image_bytes) / (1024*1024):.1f}MB > {max_size_mb}MB")
            report['is_safe'] = False
            raise ImageSanitizationError(f"文件大小超过限制: {max_size_mb}MB")

    def _validate_file_header(self, image_bytes: bytes, report: Dict[str, Any]) -> None:
        """验证文件头，检测可执行文件伪装"""
        # 检查危险文件签名
        for signature in self.dangerous_signatures:
            if image_bytes.startswith(signature):
                signature_hex = signature.hex()
                report['security_issues'].append(f"检测到危险文件签名: {signature_hex}")
                report['is_safe'] = False
                raise ImageSanitizationError("检测到伪装的恶意文件")

        # 使用python-magic进行深度文件类型检测
        try:
            detected_type = magic.from_buffer(image_bytes[:2048], mime=True)
            report['format_detected'] = detected_type

            # 检查MIME类型
            allowed_mimes = {'image/jpeg', 'image/png', 'image/bmp', 'image/x-ms-bmp'}
            if detected_type not in allowed_mimes:
                report['security_issues'].append(f"不支持的文件类型: {detected_type}")
                report['is_safe'] = False
                raise ImageSanitizationError(f"不支持的文件类型: {detected_type}")

        except Exception as e:
            report['security_issues'].append(f"文件类型检测失败: {e}")
            # 不直接失败，但标记为潜在风险

    def _scan_for_malicious_payloads(self, image_bytes: bytes, report: Dict[str, Any]) -> None:
        """扫描恶意载荷"""
        try:
            # 将字节转换为文本进行模式匹配（仅扫描安全部分）
            text_sample = image_bytes[:1024].decode('utf-8', errors='ignore')

            for pattern in self.malicious_patterns:
                if re.search(pattern, text_sample, re.IGNORECASE | re.DOTALL):
                    report['security_issues'].append(f"检测到恶意模式: {pattern}")
                    report['is_safe'] = False
                    raise ImageSanitizationError("检测到恶意载荷")

        except UnicodeDecodeError:
            # 无法解码为文本，这在图像文件中是正常的
            pass
        except Exception as e:
            logger.warning("恶意载荷扫描失败", error=str(e))

    def _load_and_validate_image(self, image_bytes: bytes, report: Dict[str, Any]) -> Image.Image:
        """加载并验证图像"""
        try:
            with Image.open(io.BytesIO(image_bytes)) as img:
                # 验证图像
                img.verify()  # 验证图像数据完整性

                # 重新打开（verify()会关闭文件）
                img = Image.open(io.BytesIO(image_bytes))

                # 检查图像格式
                if img.format not in self.allowed_formats:
                    report['security_issues'].append(f"不支持的图像格式: {img.format}")
                    raise ImageSanitizationError(f"不支持的图像格式: {img.format}")

                # 检查图像尺寸
                width, height = img.size
                if width < 10 or height < 10:
                    report['security_issues'].append(f"图像尺寸过小: {width}x{height}")
                    raise ImageSanitizationError(f"图像尺寸过小: {width}x{height}")

                if width > 8192 or height > 8192:
                    report['security_issues'].append(f"图像尺寸过大: {width}x{height}")
                    raise ImageSanitizationError(f"图像尺寸过大: {width}x{height}")

                # 检查像素数量
                total_pixels = width * height
                if total_pixels > 67108864:  # 8192*8192
                    report['security_issues'].append(f"像素数过多: {total_pixels:,}")
                    raise ImageSanitizationError(f"图像像素数过多: {total_pixels:,}")

                # 检查颜色模式
                if img.mode not in ['RGB', 'RGBA', 'L', 'P']:
                    report['cleaning_actions'].append(f"转换颜色模式: {img.mode} -> RGB")
                    img = img.convert('RGB')

                report['cleaning_actions'].append(f"图像验证通过: {width}x{height}, 格式: {img.format}")
                return img

        except UnidentifiedImageError as e:
            report['security_issues'].append(f"无法识别的图像格式: {e}")
            raise ImageSanitizationError(f"无效的图像文件: {e}")
        except Exception as e:
            report['security_issues'].append(f"图像加载失败: {e}")
            raise ImageSanitizationError(f"图像加载失败: {e}")

    def _force_reencoding(self, image: Image.Image, report: Dict[str, Any]) -> bytes:
        """强制重新编码，剥离所有非像素数据"""
        try:
            # 确保为RGB格式
            if image.mode != 'RGB':
                image = image.convert('RGB')
                report['cleaning_actions'].append("转换为RGB格式")

            # 创建输出缓冲区
            output_buffer = io.BytesIO()

            # 使用PNG格式重新保存（无损，压缩非像素数据）
            image.save(output_buffer, format='PNG', optimize=False)

            # 获取净化后的数据
            sanitized_bytes = output_buffer.getvalue()
            output_buffer.close()

            report['cleaning_actions'].append("强制PNG重新编码完成")
            report['metadata_removed'] = {
                'exif_data': '已移除',
                'color_profile': '已标准化',
                'metadata': '已清理'
            }

            return sanitized_bytes

        except Exception as e:
            raise ImageSanitizationError(f"图像重新编码失败: {e}")

    def _final_validation(self, sanitized_bytes: bytes, report: Dict[str, Any]) -> None:
        """最终验证净化结果"""
        try:
            # 验证净化后的文件是否为有效图像
            with Image.open(io.BytesIO(sanitized_bytes)) as img:
                img.verify()

            # 检查文件大小变化
            size_ratio = len(sanitized_bytes) / report['original_size']
            if size_ratio > 5.0:  # 文件大小异常增长
                report['security_issues'].append(f"净化后文件大小异常增长: {size_ratio:.1f}x")
                report['is_safe'] = False

            # 最终安全确认
            if report['is_safe'] and not report['security_issues']:
                report['cleaning_actions'].append("最终验证通过")
            else:
                report['cleaning_actions'].append("最终验证发现潜在风险")

        except Exception as e:
            report['security_issues'].append(f"最终验证失败: {e}")
            report['is_safe'] = False


class PolyglotDetector:
    """Polyglot文件检测器"""

    def __init__(self):
        # Polyglot文件特征
        self.polyglot_indicators = [
            # 多重文件头
            (b'\x89PNG\r\n\x1a\n', b'GIF8'),  # PNG + GIF
            (b'\xff\xd8\xff', b'\x89PNG'),     # JPEG + PNG
            (b'BM', b'GIF8'),                  # BMP + GIF
            (b'PK\x03\x04', b'\xff\xd8\xff'),   # ZIP + JPEG
        ]

    def detect_polyglot(self, data: bytes) -> bool:
        """检测是否为Polyglot文件"""
        try:
            # 检查是否包含多个文件头
            headers_found = []

            # 常见图像文件头
            image_headers = [
                b'\xff\xd8\xff\xe0',  # JPEG
                b'\xff\xd8\xff\xe1',  # JPEG
                b'\x89PNG\r\n\x1a\n',  # PNG
                b'GIF87a',            # GIF87a
                b'GIF89a',            # GIF89a
                b'BM',                # BMP
                b'MZ',                # PE
                b'\x7fELF',            # ELF
            ]

            for header in image_headers:
                if data.startswith(header):
                    headers_found.append(header)
                elif header in data[:1000]:  # 在前1KB中搜索
                    headers_found.append(header)

            # 如果找到多个不同的文件头，可能是Polyglot文件
            unique_headers = set(headers_found)
            if len(unique_headers) > 1:
                logger.warning("检测到Polyglot文件特征", headers_found=[h.hex() for h in unique_headers])
                return True

            # 检查文件结构异常
            if self._has_suspicious_structure(data):
                return True

            return False

        except Exception as e:
            logger.error("Polyglot检测失败", error=str(e))
            return False

    def _has_suspicious_structure(self, data: bytes) -> bool:
        """检查可疑的文件结构"""
        try:
            # 检查是否包含可执行代码的迹象
            suspicious_strings = [
                b'eval(',
                b'exec(',
                b'system(',
                b'shell_exec',
                b'base64_decode',
                b'<script',
                b'javascript:',
                b'data:text/html',
            ]

            for suss_str in suspicious_strings:
                if suss_str in data.lower():
                    return True

            # 检查文件结尾是否异常
            if len(data) > 100:
                tail = data[-100:]
                if tail.count(b'\x00') > 50:  # 大量空字节
                    return True

            return False

        except Exception:
            return False


# 全局实例
_image_sanitizer = ImageSanitizer()
_polyglot_detector = PolyglotDetector()


def sanitize_image_upload(image_bytes: bytes, max_size_mb: int = 50) -> Tuple[bytes, Dict[str, Any]]:
    """
    净化上传图像的便捷函数

    Args:
        image_bytes: 图像字节数据
        max_size_mb: 最大文件大小（MB）

    Returns:
        (净化后的图像字节数据, 处理报告)
    """
    return _image_sanitizer.sanitize_image(image_bytes, max_size_mb)


def detect_polyglot_attack(image_bytes: bytes) -> bool:
    """
    检测Polyglot攻击的便捷函数

    Args:
        image_bytes: 图像字节数据

    Returns:
        是否为Polyglot文件
    """
    return _polyglot_detector.detect_polyglot(image_bytes)


# 导出主要类和函数
__all__ = [
    'ImageSanitizationError',
    'ImageSanitizer',
    'PolyglotDetector',
    'sanitize_image_upload',
    'detect_polyglot_attack'
]